Skip to content

Update algorithmic choices for soap and clean up code #68

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
skyw merged 13 commits into from
Nov 5, 2025
Merged

Update algorithmic choices for soap and clean up code #68

skyw merged 13 commits into from
Nov 5, 2025

Conversation

@skyw
Copy link
Contributor

@skyw skyw commented Oct 31, 2025

Couple of changes (subject to change further before out of Draft):

  • Update is not skipped the first step out of adam warmup
  • Current gradient is used in preconditioning
  • First Kronecker factor is updated with Shamoo beta
  • Explicitly differentiate 0-based and 1-based step count for different purposes.

Also pulled in https://github.com/nikhilvyas/SOAP as testing reference. Code is shuffled to match updated algorithmic choices. A version matching unmodified (almost) reference is archived in branch https://github.com/NVIDIA-NeMo/Emerging-Optimizers/tree/skyw/soap_vs_reference_backup_donot_merge

skyw added 10 commits October 30, 2025 08:56
@skyw
remove rarely used functions in soap
39d6db6 
Signed-off-by: Hao Wu <[email protected]>
@skyw
update adam tests
b3dea3d 
Signed-off-by: Hao Wu <[email protected]>
@skyw
Add soap reference test
eaa2d32 
Signed-off-by: Hao Wu <[email protected]>
@skyw
add test to compare update with reference
5ebae6a 
Signed-off-by: Hao Wu <[email protected]>
@skyw
add more test to match reference
a4e3006 
Signed-off-by: Hao Wu <[email protected]>
@skyw
shuffle soap and reference to match each other
8180404 
Signed-off-by: Hao Wu <[email protected]>
@skyw
change seed to per test class
e854332 
Signed-off-by: Hao Wu <[email protected]>
@skyw
Clean up code
857325c 
Signed-off-by: Hao Wu <[email protected]>
@skyw
Update some test shmoo
41b170c 
Signed-off-by: Hao Wu <[email protected]>
@skyw
update some magical number
2b05036 
Signed-off-by: Hao Wu <[email protected]>
@skyw skyw requested a review from mkhona-nvidia October 31, 2025 21:41
@copy-pr-bot
Copy link

copy-pr-bot bot commented Oct 31, 2025

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@skyw
Copy link
Contributor Author

skyw commented Oct 31, 2025

/ok to test 2b05036

mkhona-nvidia
mkhona-nvidia reviewed Nov 3, 2025
View reviewed changes
mkhona-nvidia
mkhona-nvidia reviewed Nov 3, 2025
View reviewed changes
mkhona-nvidia
mkhona-nvidia reviewed Nov 3, 2025
View reviewed changes
Copy link
Contributor

@mkhona-nvidia mkhona-nvidia left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

soap.py LGTM

skyw
skyw commented Nov 3, 2025
View reviewed changes
@skyw
Copy link
Contributor Author

skyw commented Nov 4, 2025

/ok to test e302bf7

mkhona-nvidia
mkhona-nvidia reviewed Nov 5, 2025
View reviewed changes
mkhona-nvidia
mkhona-nvidia approved these changes Nov 5, 2025
View reviewed changes
Copy link
Contributor

@mkhona-nvidia mkhona-nvidia left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

@skyw
Copy link
Contributor Author

skyw commented Nov 5, 2025

/ok to test 53cb513

@skyw skyw marked this pull request as ready for review November 5, 2025 17:54
@skyw skyw enabled auto-merge (squash) November 5, 2025 17:54
@skyw skyw merged commit 16f8399 into main Nov 5, 2025
21 of 23 checks passed
@skyw skyw deleted the skyw/soap_further_cleanup branch November 5, 2025 18:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@mkhona-nvidia mkhona-nvidia mkhona-nvidia approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@skyw @mkhona-nvidia